package t20240514_lineage;

import com.alibaba.fastjson.JSONObject;
import io.debezium.util.Strings;
import org.apache.calcite.plan.RelOptTable;
import org.apache.calcite.rel.RelNode;
import org.apache.calcite.rel.metadata.JaninoRelMetadataProvider;
import org.apache.calcite.rel.metadata.RelColumnOrigin;
import org.apache.calcite.rel.metadata.RelMetadataQuery;
import org.apache.calcite.rel.metadata.RelMetadataQueryBase;
import org.apache.calcite.rel.type.RelDataTypeField;
import org.apache.commons.collections.CollectionUtils;
import org.apache.flink.streaming.api.environment.StreamExecutionEnvironment;
import org.apache.flink.table.api.bridge.java.StreamTableEnvironment;
import org.apache.flink.table.api.bridge.java.internal.StreamTableEnvironmentImpl;
import org.apache.flink.table.operations.CatalogSinkModifyOperation;
import org.apache.flink.table.operations.Operation;
import org.apache.flink.table.operations.ddl.CreateTableOperation;
import org.apache.flink.table.plan.cost.FlinkDefaultRelMetadataProvider;
import org.apache.flink.table.planner.operations.PlannerQueryOperation;
import org.apache.flink.table.planner.plan.schema.TableSourceTable;

import java.util.*;

public class Test1_1 {

    public static final String SQL_EMPTY_STR = "[\\s\\t\\n\\r]";



    private static final String CTAS_REGEX = "CREATE\\s+TABLE(?s).*AS\\s+SELECT(?s).*";

    public static void main(String[] args) {



        StreamExecutionEnvironment env = StreamExecutionEnvironment.getExecutionEnvironment();
        StreamTableEnvironment tbEnv = StreamTableEnvironment.create(env);


        String sql1 = "create table o_scdn_vod_xserver_ngx_httpcode (\n" +
                "  `stime` string,\n" +
                "  `type`  string,\n" +
                "  `hostname`  string,\n" +
                "  `domain`  string,\n" +
                "  `status`  int,\n" +
                "  `status_cnt`  int,\n" +
                "  `port`  int, \n" +
                "  `level` string,\n" +
                "  rowtime as to_timestamp(`stime`),\n" +
                "  proctime as proctime(),\n" +
                "  watermark for rowtime as rowtime - interval '20' second\n" +
                ") with (\n" +
                "  'properties.bootstrap.servers' = 'bdyun-suzhou-tw04a0015.su.baidu.internal:6667,bdyun-suzhou-tw04b0023.su.baidu.internal:6667,bdyun-suzhou-tw04b0025.su.baidu.internal:6667,bdyun-suzhou-tw04d0583.su.baidu.internal:6667,bdyun-suzhou-tw04d0581.su.baidu.internal:6667',\n" +
                "  'connector' = 'kafka',\n" +
                "  'format' = 'avro-event',\n" +
                "  'topic' = 'o_scdn_vod_xserver_ngx_httpcode',\n" +
                "  'properties.group.id' = 'flink-lichen-o_scdn_vod_xserver_ngx_httpcode-lc-kafka-scdn2',\n" +
                "  'properties.security.protocol' = 'SASL_PLAINTEXT',\n" +
                "  'scan.startup.mode' = 'latest-offset',\n" +
                "  'properties.sasl.kerberos.service.name' = 'kafka'\n" +
                ")";
        String sql2 = "create table t_domain_of_data_platform (\n" +
                "  `userid`  bigint,\n" +
                "  `domain_name` string\n" +
                ") with (\n" +
                "  'connector' = 'tidb',\n" +
                "  'driver' = 'com.mysql.jdbc.Driver',\n" +
                "  'lookup.cache.ttl' = '30min',\n" +
                "  'table-name' = 't_domain_of_data_platform',\n" +
                "  'url' = 'jdbc:mysql://tidb-lb.data.onething.internal:4000/dim?useUnicode=true&useSSL=false&characterEncoding=UTF-8&allowMultiQueries=true&autoReconnect=true',\n" +
                "  'lookup.max-retries' = '3',\n" +
                "  'username' = 'beta',\n" +
                "    'password' = 'beta@9898w',\n" +
                "  'query-sql' = 'select account_id as userid, domain_name from t_domain_of_data_platform'\n" +
                ")";
        String sql3 = "create table app_i1_scdn_vod_xserver_httpcode_stat (\n" +
                "  `stime` timestamp,\n" +
                "  `grpid` int,\n" +
                "  `account_id`  string,\n" +
                "  `domain`  string,\n" +
                "  `hostname`  string,\n" +
                "  `level` string,\n" +
                "  `status`  int,\n" +
                "  `status_cnt_total`  int\n" +
                ") with (\n" +
                "  'database' = 'scdn_app_lo',\n" +
                "  'password' = 'rNXB>=dQCe3Ri-CGg*B?',\n" +
                "  'table-mode' = 'single',\n" +
                "  'sink.write-mode' = 'local',\n" +
                "  'connector' = 'clickhouse',\n" +
                "  'sink.flush-interval' = '10s',\n" +
                "  'sink.batch-size' = '10000',\n" +
                "  'url' = 'clickhouse://ck-lb.data.onething.internal:8123',\n" +
                "  'table' = 'app_i1_scdn_vod_xserver_httpcode_stat',\n" +
                "  'username' = 'spark'\n" +
                ")";
        String sql4 = "create view o_scdn_vod_xserver_ngx_httpcode_fliter as\n" +
                "select\n" +
                "    *\n" +
                "from o_scdn_vod_xserver_ngx_httpcode\n" +
                "where\n" +
                "    (\n" +
                "    (level = 'edge' and `port` in (80, 443, 9102, 9305))\n" +
                "    or (level = 'area' and `port` in (80, 9102))\n" +
                "    or (level = 'root' and `port` = 80)\n" +
                "    )\n" +
                "    and `domain` <> '\\N'\n" +
                "    and `hostname` <> '\\N'\n" +
                "    and level <> '\\N'";

        String sql5 = "create view o_scdn_vod_xserver_ngx_httpcode_fliter_join_rules as\n" +
                "select\n" +
                "    ori.`hostname`,\n" +
                "    ori.`domain` as field,\n" +
                "    ori.`domain`,\n" +
                "    ori.`status`,\n" +
                "    ori.`level`,\n" +
                "    ori.`status_cnt`,\n" +
                "    ori.`rowtime`,\n" +
                "    proctime() as proctime\n" +
                "from o_scdn_vod_xserver_ngx_httpcode_fliter ori";

        String sql6 =
                "create view o_scdn_vod_xserver_ngx_httpcode_fliter_join_rules_plat as\n" +
                "select\n" +
                "    ifnull(split_index(oria.`field`, '|', 1), cast(plat.userid as string)) `account_id`,\n" +
                "    oria.`hostname`,\n" +
                "    split_index(oria.`field`, '|', 0) as dom,\n" +
                "    oria.`status`,\n" +
                "    oria.`level`,\n" +
                "    oria.`status_cnt`,\n" +
                "    oria.`rowtime`\n" +
                "from o_scdn_vod_xserver_ngx_httpcode_fliter_join_rules oria\n" +
                "left join t_domain_of_data_platform for system_time as of oria.proctime plat on split_index(oria.`field`, '|', 0) = plat.`domain_name`";

        String sql7 =
                "create view o_scdn_vod_xserver_ngx_httpcode_fliter_join_rules_plat_group as\n" +
                "select\n" +
                "    window_start,\n" +
                "    window_end,\n" +
                "    cast(`account_id` as string) `account_id`,\n" +
                "    `hostname`,\n" +
                "    `dom`,\n" +
                "    `status`,\n" +
                "    `level`,\n" +
                "    sum(`status_cnt`) as `status_sum`\n" +
                "from table(tumble(table o_scdn_vod_xserver_ngx_httpcode_fliter_join_rules_plat, descriptor(`rowtime`), interval '60' second))\n" +
                "group by `account_id` ,`hostname`, `dom`, `status`, `level`, window_start, window_end";
        String sql8 = "create view xxxxxx as\n" +
                "select\n" +
                "    window_start as stime,\n" +
                "    1 `grpid`,\n" +
                "    case `account_id` when 'null' then null else `account_id` end as `account_id`,\n" +
                "    `dom` as `domain`,\n" +
                "    '-9999' as `hostname`,\n" +
                "    `level`,\n" +
                "    `status`,\n" +
                "    sum(`status_sum`) as `status_cnt_total`\n" +
                "from o_scdn_vod_xserver_ngx_httpcode_fliter_join_rules_plat_group\n" +
                "group by  case `account_id` when 'null' then null else `account_id` end , `dom`, `status`, `level`, window_start, window_end\n" +
                "union all\n" +
                "select\n" +
                "    window_start as stime,\n" +
                "    2 `grpid`,\n" +
                "    case `account_id` when 'null' then null else `account_id` end as `account_id`,\n" +
                "    `dom` as `domain`,\n" +
                "    `hostname`,\n" +
                "    `level`,\n" +
                "    `status`,\n" +
                "    sum(`status_sum`) as `status_cnt_total`\n" +
                "from o_scdn_vod_xserver_ngx_httpcode_fliter_join_rules_plat_group\n" +
                "group by case `account_id` when 'null' then null else `account_id` end , `dom`,`hostname`,`level`,`status`, window_start, window_end";


        String sql9 = "insert into app_i1_scdn_vod_xserver_httpcode_stat select * from xxxxxx";

        parseSql(tbEnv, sql1);
        parseSql(tbEnv, sql2);
        parseSql(tbEnv, sql3);
        parseSql(tbEnv, sql4);
        parseSql(tbEnv, sql5);
        parseSql(tbEnv, sql6);
        parseSql(tbEnv, sql7);
        parseSql(tbEnv, sql8);
        parseSql(tbEnv, sql9);


    }

    private static void parseSql(StreamTableEnvironment tbEnv, String singleSql) {
        SqlType operationType = getOperationType(singleSql);
        StreamTableEnvironmentImpl tbEnvImpl = (StreamTableEnvironmentImpl) tbEnv;
        if (operationType.equals(SqlType.INSERT)) {

            RelMetadataQueryBase.THREAD_PROVIDERS
                    .set(JaninoRelMetadataProvider.of(FlinkDefaultRelMetadataProvider.INSTANCE()));
            System.out.println("singleSql :" + singleSql);
            List<Operation> list = tbEnvImpl.getParser().parse(singleSql);

            Operation operation = list.get(0);
            System.out.println(operation);



            CatalogSinkModifyOperation sinkOperation = (CatalogSinkModifyOperation) operation;

            PlannerQueryOperation queryOperation = (PlannerQueryOperation) sinkOperation.getChild();
            RelNode relNode = queryOperation.getCalciteTree();

            String sinkTbName = sinkOperation.getTableIdentifier().asSummaryString();
            System.out.println("sinkTbName:" + sinkTbName);
            System.out.println("relNode:" + relNode);


            System.out.println("----------------------------------------------------------");


            List<String> targetColumnList = tbEnv.from(sinkTbName).getResolvedSchema().getColumnNames();
            List<RelDataTypeField> fieldList = relNode.getRowType().getFieldList();

            System.out.println("targetColumnList:" + targetColumnList);
            System.out.println("sourceColumnList:" + fieldList);

            System.out.println("----------------------------------------------------------");

            RelMetadataQuery metadataQuery = relNode.getCluster().getMetadataQuery();
            List<LineageRel> resultList = new ArrayList<>();


            System.out.println();

            for (int i = 0; i < targetColumnList.size(); i++) {

                String targetColumn = targetColumnList.get(i);

                Set<RelColumnOrigin> columnOrigins = metadataQuery.getColumnOrigins(relNode, i);
//                System.out.println(targetColumnList.get(i) + " source:" + columnOrigins.size());

                if (CollectionUtils.isNotEmpty(columnOrigins)) {
                    for (RelColumnOrigin columnOrigin : columnOrigins) {

                        RelOptTable originTable = columnOrigin.getOriginTable();
                        String tableName = Strings.join(".", originTable.getQualifiedName());


                        int ordinal = columnOrigin.getOriginColumnOrdinal();

                        TableSourceTable tableSourceTable = (TableSourceTable) originTable;
                        List<String> columnNames = tableSourceTable.catalogTable().getResolvedSchema().getColumnNames();
                        String sourceColumn = columnNames.get(ordinal);

                        Map<String, String> options = tableSourceTable.catalogTable().getOptions();
                        String sourceConnector = options.get("connector");


                        // add record
                        resultList.add(LineageRel.build(
                                tableName, sourceColumn,
                                sinkTbName, targetColumn,
                                null));

                    }
                } else {
                    System.err.println("not found:" + targetColumn);
                }
            }


            System.out.println("resultList:" + resultList.size());


            LineageResult lineageResult = buildLineage(resultList);
            JSONObject jobj = (JSONObject)JSONObject.toJSON(lineageResult);
            System.out.println(jobj.toJSONString());

        } else if (!operationType.equals(SqlType.SELECT) && !operationType.equals(SqlType.PRINT)) {
            tbEnv.executeSql(singleSql);

            if (operationType.equals(SqlType.CREATE)) {
                Operation operation = tbEnvImpl.getParser().parse(singleSql).get(0);
                if (operation instanceof CreateTableOperation) {
                    CreateTableOperation createTableOperation = (CreateTableOperation) operation;
                    String connector = createTableOperation.getCatalogTable().getOptions().get("connector");
                    System.out.println("option:" + createTableOperation.getTableIdentifier().asSummaryString() + " " + connector);
                }

            }



        }


    }



    public static SqlType getOperationType(String sql) {
        String sqlTrim = sql.replaceAll(SQL_EMPTY_STR, " ").trim().toUpperCase();
        return Arrays.stream(SqlType.values())
                .filter(sqlType -> sqlType.match(sqlTrim))
                .findFirst()
                .orElse(SqlType.UNKNOWN);
    }




    public static LineageResult buildLineage(List<LineageRel> lineageRelList) {
        List<LineageRelation> relations = new ArrayList<>();
        Map<String, LineageTable> tableMap = new HashMap<>();
        int tableIndex = 1;
        int relIndex = 1;
        for (LineageRel lineageRel : lineageRelList) {
            String sourceTablePath = lineageRel.getSourceTablePath();
            String sourceTableId = null;
            String targetTableId = null;
            if (tableMap.containsKey(sourceTablePath)) {
                LineageTable lineageTable = tableMap.get(sourceTablePath);
                LineageColumn lineageColumn =
                        LineageColumn.build(lineageRel.getSourceColumn(), lineageRel.getSourceColumn());
                if (!lineageTable.getColumns().contains(lineageColumn)) {
                    lineageTable.getColumns().add(lineageColumn);
                }
                sourceTableId = lineageTable.getId();
            } else {
                tableIndex++;
                LineageTable lineageTable = LineageTable.build(String.valueOf(tableIndex), sourceTablePath, "XXX");
                lineageTable
                        .getColumns()
                        .add(LineageColumn.build(lineageRel.getSourceColumn(), lineageRel.getSourceColumn()));
                tableMap.put(sourceTablePath, lineageTable);
                sourceTableId = lineageTable.getId();
            }
            String targetTablePath = lineageRel.getTargetTablePath();
            if (tableMap.containsKey(targetTablePath)) {
                LineageTable lineageTable = tableMap.get(targetTablePath);
                LineageColumn lineageColumn =
                        LineageColumn.build(lineageRel.getTargetColumn(), lineageRel.getTargetColumn());
                if (!lineageTable.getColumns().contains(lineageColumn)) {
                    lineageTable.getColumns().add(lineageColumn);
                }
                targetTableId = lineageTable.getId();
            } else {
                tableIndex++;
                LineageTable lineageTable = LineageTable.build(String.valueOf(tableIndex), targetTablePath, "XXX");
                lineageTable
                        .getColumns()
                        .add(LineageColumn.build(lineageRel.getTargetColumn(), lineageRel.getTargetColumn()));
                tableMap.put(targetTablePath, lineageTable);
                targetTableId = lineageTable.getId();
            }
            LineageRelation lineageRelation = LineageRelation.build(
                    sourceTableId, targetTableId, lineageRel.getSourceColumn(), lineageRel.getTargetColumn());
            if (!relations.contains(lineageRelation)) {
                relIndex++;
                lineageRelation.setId(relIndex + "");
                relations.add(lineageRelation);
            }
        }
        List<LineageTable> tables = new ArrayList<LineageTable>(tableMap.values());
        return LineageResult.build(tables, relations);
    }
}
